// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/common/base/exceptions.h>
#include <pollux/exec/window_function.h>
#include <pollux/expression/function_signature.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::window::prestosql {
    namespace {
        class CumeDistFunction : public exec::WindowFunction {
        public:
            explicit CumeDistFunction() : WindowFunction(DOUBLE(), nullptr, nullptr) {
            }

            void resetPartition(const exec::WindowPartition *partition) override {
                runningTotal_ = 0;
                cumeDist_ = 0;
                currentPeerGroupStart_ = -1;
                numPartitionRows_ = partition->numRows();
            }

            void apply(
                const BufferPtr &peerGroupStarts,
                const BufferPtr &peerGroupEnds,
                const BufferPtr & /*frameStarts*/,
                const BufferPtr & /*frameEnds*/,
                const SelectivityVector &validRows,
                vector_size_t resultOffset,
                const VectorPtr &result) override {
                int numRows = peerGroupStarts->size() / sizeof(vector_size_t);
                auto *rawValues = result->as_flat_vector<double>()->mutableRawValues();
                auto *peerGroupStartsVector = peerGroupStarts->as<vector_size_t>();
                auto *peerGroupEndsVector = peerGroupEnds->as<vector_size_t>();

                for (int i = 0; i < numRows; i++) {
                    auto peerStart = peerGroupStartsVector[i];
                    if (peerStart != currentPeerGroupStart_) {
                        currentPeerGroupStart_ = peerStart;
                        runningTotal_ += peerGroupEndsVector[i] - peerStart + 1;
                        cumeDist_ = static_cast<double>(runningTotal_) / numPartitionRows_;
                    }
                    rawValues[resultOffset + i] = cumeDist_;
                }
            }

        private:
            int64_t runningTotal_ = 0;
            double cumeDist_ = 0;
            int64_t currentPeerGroupStart_ = -1;
            vector_size_t numPartitionRows_ = 1;
        };
    } // namespace

    void registerCumeDist(const std::string &name) {
        std::vector<exec::FunctionSignaturePtr> signatures{
            exec::FunctionSignatureBuilder().returnType("double").build(),
        };

        exec::registerWindowFunction(
            name,
            std::move(signatures),
            exec::WindowFunction::Metadata::defaultMetadata(),
            [name](
        const std::vector<exec::WindowFunctionArg> & /*args*/,
        const TypePtr & /*resultType*/,
        bool /*ignoreNulls*/,
        pollux::memory::MemoryPool * /*pool*/,
        HashStringAllocator * /*stringAllocator*/,
        const pollux::core::QueryConfig & /*queryConfig*/)
        -> std::unique_ptr<exec::WindowFunction> {
                return std::make_unique<CumeDistFunction>();
            });
    }
} // namespace kumo::pollux::window::prestosql
